# 用于统计分割·标志所能分割出句子的数量，确保分割出的EEG和文本对的上
import numpy as np
import scipy
import os
import pandas as pd
import mat73
import pickle

eeg, text = [], []
for task_name in os.listdir("./rawdata/zuco2/"):
    for sbj_name in os.listdir(f"./rawdata/zuco2/{task_name}/Raw data"):

        file_name = f"./rawdata/zuco2/{task_name}/Matlab files/results{sbj_name}_NR.mat"
        text_ = mat73.loadmat(file_name)['sentenceData']['content']

        folder = F"./rawdata/zuco2/{task_name}/Raw data/{sbj_name}/"
        filename_li = [f"{sbj_name}_NR{_}_EEG.mat" for _ in range(1, 8)]
        eeg_seg = []
        for file_name in filename_li:
            data = scipy.io.loadmat(folder + file_name, squeeze_me=True, struct_as_record=False)
            raweeg = data["EEG"].data
            if "EEG" in data.keys():
                event = data["EEG"].event
            else:
                event = data["event"]
            
            
            cnt_10, cnt_11, cnt_12, cnt_13 = None, None, None, None
            cnt = 0
            for each in event:
                if each.value == "trigger":
                    if each.type.strip() == "10":
                        cnt_10 = each.latency
                    elif each.type.strip() == "11":
                        cnt_11 = each.latency
                        eeg_seg.append(raweeg[:, cnt_10:cnt_11+1])
                        cnt_10, cnt_11 = None, None
                        cnt += 1
                    elif each.type.strip() == "12":
                        cnt_12 = each.latency
                    elif each.type.strip() == "13":
                        cnt_13 = each.latency
                        eeg_seg.append(raweeg[:, cnt_12:cnt_13+1])
                        cnt_12, cnt_13 = None, None
                        cnt += 1
            assert len(eeg_seg) == len(text_)
        eeg += eeg_seg
        text += text_
assert len(eeg) == len(text)
dic = {"eeg": eeg, "text": text}
pickle.dump(dic, open("./data/zuco2.pkl", "wb"))


# folder = "/data2/npc/lixiang/playground/aaa/mind-whisper/rawdata/zuco2/task1 - NR/Raw data/YAK/"
# for file_name in os.listdir(folder):
#     if "EEG" not in file_name:
#         continue
#     data = scipy.io.loadmat(folder + file_name, squeeze_me=True, struct_as_record=False)
#     if "EEG" in data.keys():
#         event = data["EEG"].event
#     else:
#         event = data["event"]
#     cnt_10, cnt_11, cnt_12, cnt_13 = 0, 0, 0, 0
#     for each in event:
#         if each.value == "trigger":
#             if each.type.strip() == "10":
#                 cnt_10 += 1
#             elif each.type.strip() == "11":
#                 cnt_11 += 1
#             elif each.type.strip() == "12":
#                 cnt_12 += 1
#             elif each.type.strip() == "13":
#                 cnt_13 += 1
#     print(f"file_name: {file_name}")
#     print(f"cnt_10: {cnt_10}, cnt_11: {cnt_11}, cnt_12: {cnt_12}, , cnt_13: {cnt_13}")
# # ------------只数单个文件的分割标志--------------------------------------------------------------------
# file_name = "/data2/npc/lixiang/playground/aaa/mind-whisper/rawdata/zuco2/task1 - NR/Raw data/YAK/YAK_NR1_EEG.mat"
# data = scipy.io.loadmat(file_name, squeeze_me=True, struct_as_record=False)
# raweeg = data['EEG'].data
# event = data['EEG'].event


# seg = {}
# for each in event:
#     if each.value=='trigger':
#         seg[each.latency] = each.type.strip()

# count_10 = sum(value == '10' for value in seg.values())
# count_11 = sum(value == '11' for value in seg.values())
# count_12 = sum(value == '12' for value in seg.values())
# count_13 = sum(value == '13' for value in seg.values())

# print(seg)
# print(f"count_10:{count_10}, count_11:{count_11}")
# print(f"count_12:{count_12}, count_13:{count_13}")
